from datetime import datetime
import os
import logging
import shutil
from util_tool import get_vmrules, run_command
from info_config import (
    kubeconfig_path,
    vmrule_namespace,
    backup_path,
    kube_command_path,
    access_key_id,
    access_key_secret,
    backup_bucket_name,
    backup_in_oss_path,
    oss_region_url,
)
from oss2 import Auth, Bucket

# Set up logging
script_directory = os.path.dirname(os.path.abspath(__file__))
log_file_path = os.path.join(script_directory, "backup_execution.log")
logging.basicConfig(
    filename=log_file_path,
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

def mkdir_backup(backup_directory):
    try:
        os.makedirs(backup_directory, exist_ok=True)
        logging.info(f"Directory {backup_directory} created successfully")
    except OSError as error:
        logging.error(f"Error creating directory {backup_directory}: {error}")


def backup_vmrules(vmrule_name, backup_directory):
    backup_rule_command = (
        f"{kube_command_path} get VMRule {vmrule_name} "
        f"--kubeconfig={kubeconfig_path} -n {vmrule_namespace} -o yaml > {backup_directory}/{vmrule_name}.yaml"
    )
    r = run_command(backup_rule_command)


def compress_backup_directory(backup_directory):
    compressed_file_path = f"{backup_directory}.tar.gz"
    try:
        shutil.make_archive(backup_directory, "gztar", backup_directory)
        logging.info(f"Backup directory compressed to {compressed_file_path}")
    except Exception as e:
        logging.error(f"Error compressing backup directory: {e}")
    return compressed_file_path


def upload_to_oss(file_path, oss_path):
    auth = Auth(access_key_id, access_key_secret)
    bucket = Bucket(auth, oss_region_url, backup_bucket_name)

    oss_file_path = f"{backup_in_oss_path}/{os.path.basename(file_path)}"
    try:
        bucket.put_object_from_file(oss_file_path, file_path)
        logging.info(f"File {file_path} uploaded to OSS at {oss_file_path}")
    except Exception as e:
        logging.error(f"Error uploading file to OSS: {e}")


def delete_backup_directory(backup_directory):
    try:
        shutil.rmtree(backup_directory)
        logging.info(f"Backup directory {backup_directory} deleted successfully")
    except Exception as e:
        logging.error(f"Error deleting backup directory {backup_directory}: {e}")


def main():
    # Get the current time
    now = datetime.now()
    time_string = now.strftime("%Y%m%d%H%M%S")
    backup_directory = f"{backup_path}/{time_string}"
    mkdir_backup(backup_directory)

    # Backup each VMRule
    vmrules = get_vmrules(vmrule_namespace, kubeconfig_path)
    if vmrules:
        for vmrule in vmrules:
            backup_vmrules(vmrule, backup_directory)

    # Compress the backup directory
    compressed_file_path = compress_backup_directory(backup_directory)

    # Upload the compressed file to OSS
    upload_to_oss(compressed_file_path, backup_in_oss_path)

    # Delete the backup directory
    delete_backup_directory(backup_directory)


if __name__ == "__main__":
    main()
